#ifndef SD_MEANSHIFT2_HPP
#define SD_MEANSHIFT2_HPP

#include "config.hpp"
#include "rng.hpp"

namespace MT {

// -----------------------------------------------------------------------------------------
// MeanShift0 -- abstract class
// -----------------------------------------------------------------------------------------

// Finds the maxima of a density function given discrete data sampled from that function
// URL: http://en.wikipedia.org/wiki/Mean-shift
// Data is represented by a 2D CV_REAL array of N rows and D columns.
// There are N examples in a D-dimensional space.
// If weights is given, then it has N elements reresenting the weights of the examples.
// Otherwise, please make weights.empty()==true.
//
// Output is a 2D CV_REAL array containing a number of local modes
// The number of local modes may be smaller than the number of trials
template<typename _Scalar, int _Rows>
class MeanShift0
{
protected: // internal data
    int N; // number of examples
    int D; // dimensionality

	Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> weights; // weights


public: // public interface

    // initialise new meanshift problem
    virtual bool init(Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> const &data, Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> const &weights)
    {
        N = data.cols(); D = data.rows();
        this->weights = weights;
        return true;
    }

    // run meanshift multiple times and compute multiple local modes, return the number of modes
    virtual int process(Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> &output_local_modes, Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> &output_weights) { return 0; }

};




// -----------------------------------------------------------------------------------------
// MeanShiftStandard -- standard version using an anisotropic Gaussian kernel
// -----------------------------------------------------------------------------------------

// Finds the maxima of a density function given discrete data sampled from that function
// URL: http://en.wikipedia.org/wiki/Mean-shift
// Data is represented by a 2D CV_REAL array of N rows and D columns.
// There are N examples in a D-dimensional space.
// If weights is given, then it has N elements reresenting the weights of the examples.
// Otherwise, please make weights.empty()==true.
// The Gaussian importance sampling function is assumed to be anisotropic with scale s
//   K(x_i - x) = exp(-||x_i - x||_s^2/2)
// where ||x||_s = \sqrt{\sum_{d} (x_d/s_d)^2}
// n_trials is the number of trials
// max_iter is the maximum number of iterations per trial
//
// Output is a 2D CV_REAL array containing a number of local modes
// The number of local modes may be smaller than the number of trials
//
// Functions to measure distances and to compute a weighted mean can be overloaded
// so that the class works on non-Euclidean spaces
template<typename _Scalar, int _Rows>
class MeanShiftStandard : public MeanShift0<_Scalar, _Rows>
{
public: // public data
    int n_trials; // number of trials
    int max_iter; // maximum number of iteration per trial
    Eigen::Array<_Scalar, _Rows, 1> s; // scaling factors -- Gaussian sigma for each dimension

public: // some functions that can be overloaded

    // compute the distance between each column of 'scaled_data' and a given point 'pt'
    // the default is square of L2-norm Euclidean distance
    virtual bool colwise_squared_distance(Eigen::Matrix<_Scalar, _Rows, 1> const &pt, Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> &result)
    {
		result.noalias() = ((scaled_data.colwise()-pt).colwise().squaredNorm()).transpose();
        return true;
    }

    // compute the mean of a set of weighted points
    // the default is weighted Euclidean mean
    virtual bool mean(Eigen::SparseVector<_Scalar> const &weights, Eigen::Matrix<_Scalar, _Rows, 1> &output_mean, _Scalar &output_weight)
    {
        output_weight = weights.sum();
        output_mean.noalias() = scaled_data * weights;
        output_mean /= output_weight;
        return true;
    }

public: // public interface
    MeanShiftStandard(Eigen::Array<_Scalar, _Rows, 1> const &s, int n_trials=100, int max_iter=100)
        : n_trials(n_trials), max_iter(max_iter), s(s), rng(s) {}

    // initialise new meanshift problem
    bool init(Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> const &data, Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> const &weights)
    {
        MeanShift0<_Scalar, _Rows>::init(data, weights);

        int i;

        // make std==1.0
        scaled_data.resize(MeanShift0<_Scalar, _Rows>::D, MeanShift0<_Scalar, _Rows>::N);
        for(i = 0; i < MeanShift0<_Scalar, _Rows>::D; ++i)
            scaled_data.row(i) = data.row(i) / s[i];

        // random number generator
        rng = PdfRNG<_Scalar, Eigen::Dynamic>(weights);

        return true;
    }

    // run meanshift multiple times and compute multiple local modes, return the number of modes
    int process(Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> &output_local_modes, Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> &output_weights)
    {
        std::vector<Eigen::SparseVector<_Scalar> > w;
        int n = process(output_local_modes, w);

        output_weights.resize(n);
        for(int i = 0; i < n; ++i) output_weights[i] = w[i].sum();

        return n;
    }

    // run meanshift multiple times and compute multiple local modes, return the number of modes
    virtual int process(Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> &output_local_modes2, std::vector<Eigen::SparseVector<_Scalar> > &output_weights)
    {
        // clear previous results
        std::vector<Eigen::Matrix<_Scalar, _Rows, 1>, Eigen::aligned_allocator<Eigen::Matrix<_Scalar, _Rows, 1> > > output_local_modes;
        output_weights.clear();

        Eigen::Matrix<_Scalar, _Rows, 1> x;
        Eigen::SparseVector<_Scalar> w;
        int i, n;

        for(n = 0; n < n_trials; ++n)
        {
            // find a local mode
            meanshift(scaled_data.col(rng()), x, w);

            // check if this local mode has been found
            bool found = false;
            i = output_local_modes.size(); while(--i >= 0)
                if((x-output_local_modes[i]).norm() < 0.00001)
                    { found = true; break; }

            if(!found) // add the new local mode to the output array
            {
                output_local_modes.push_back(x); // create a new point
                output_weights.push_back(w);
            }
        }

        // convert points back to the original space
        i = output_local_modes.size(); while(--i >= 0)
            unscale(output_local_modes[i]);

        // convert to matrix
        i = output_local_modes.size();
        output_local_modes2.resize(MeanShift0<_Scalar, _Rows>::D, i);
        while(--i >= 0) output_local_modes2.col(i) = output_local_modes[i];

        return output_local_modes.size();
    }

protected: // internal methods

    virtual void scale(Eigen::Matrix<_Scalar, _Rows, 1> &pt) { pt.array() /= s; } // convert from input space to scaled space
    virtual void unscale(Eigen::Matrix<_Scalar, _Rows, 1> &pt) { pt.array() *= s; } // convert from scaled space to input space

    // run meanshift once, starting with the initial point, return the number of iterations
    // input and output are in the scaled space, not the original space
    virtual int meanshift(Eigen::Matrix<_Scalar, _Rows, 1> const &initial_x, Eigen::Matrix<_Scalar, _Rows, 1> &output_x, Eigen::SparseVector<_Scalar> &output_weights)
    {
        output_x.resize(MeanShift0<_Scalar, _Rows>::D, 1); // allocate a D-dimensional point

        // variables
        int i, iter, j;
        Eigen::Matrix<_Scalar, Eigen::Dynamic, 1> w;
        _Scalar sum_w;
        Eigen::Matrix<_Scalar, _Rows, 1> x(MeanShift0<_Scalar, _Rows>::D, 1);

        // initialisation
        output_x = initial_x;

        output_weights.resize(MeanShift0<_Scalar, _Rows>::N);
        for(iter = 0; iter < max_iter; ++iter)
        {
            // update output_weights
            colwise_squared_distance(output_x, w);
            for(i = 0, j = 0; i < MeanShift0<_Scalar, _Rows>::N; ++i) if(w[i] <= 100) ++j;
            if(!j) break; // can't find any nearby point
            output_weights.m_data.resize(j);
            for(i = 0, j = 0; i < MeanShift0<_Scalar, _Rows>::N; ++i) if(w[i] <= 100)
            {
                output_weights.m_data.index(j) = i;
                output_weights.m_data.value(j++) = exp(-0.5*w[i]) * MeanShift0<_Scalar, _Rows>::weights[i];
            }

            // find new mean (x)
            mean(output_weights, x, sum_w);
            if(iszero(sum_w)) break; // can't find any nearby point

            // check if new mean same as old mean
            if((x-output_x).cwiseAbs().sum() < 0.000001) break;

            // update new mean
            output_x = x;
        }

        return iter;
    }

protected: // internal data
    // random number generator
    PdfRNG<_Scalar, Eigen::Dynamic> rng;

    // scaled data of the input data so that std==1.0 in all directions
    Eigen::Matrix<_Scalar, _Rows, Eigen::Dynamic> scaled_data;
};



// -----------------------------------------------------------------------------------------
// MeanShiftFLANN -- standard version using an anisotropic Gaussian kernel and FLANN
// -----------------------------------------------------------------------------------------
// TODO: UnitTest

// Finds the maxima of a density function given discrete data sampled from that function
// URL: http://en.wikipedia.org/wiki/Mean-shift
// Data is represented by a 2D CV_REAL array of N rows and D columns.
// There are N examples in a D-dimensional space.
// If weights is given, then it has N elements reresenting the weights of the examples.
// Otherwise, please make weights.empty()==true.
// The Gaussian importance sampling function is assumed to be anisotropic with scale s
//   K(x_i - x) = exp(-||x_i - x||_s^2/2)
// where ||x||_s = \sqrt{\sum_{d} (x_d/s_d)^2}
// n_trials is the number of trials
// max_iter is the maximum number of iterations per trial
//
// Output is a 2D CV_REAL array containing a number of local modes
// The number of local modes may be smaller than the number of trials
/*
class MeanShiftFLANN : public MeanShiftStandard
{
public: // public data
    cv::flann::IndexParams const &ip; // IndexParams for FLANN

public: // public interface
    MeanShiftFLANN(cv::flann::IndexParams const &ip, cv::Mat const &s, int n_trials=100, int max_iter=100)
        : MeanShiftStandard(s, n_trials, max_iter), ip(ip) { index = 0; }

    ~MeanShiftFLANN() { if(index) delete index; }

    // initialise new meanshift problem
    bool init(cv::Mat const &data, cv::Mat const &weights);

protected: // internal methods
    // run meanshift once, starting with the initial point, return the number of iterations
    // input and output are in the scaled space, not the original space
    int meanshift(cv::Mat const &initial_x, cv::Mat &output_x, Eigen::SparseVector<_Scalar> &output_weights);

protected: // internal data
    // flann index
    cv::Mat data_float32;
    std::vector<int> indices;
    std::vector<float32> dists;
    cv::flann::SearchParams search_params;
    std::vector<float32> query;
    cv::flann::Index *index;
};
*/

// ****************************************************************************************
// MeanShiftFLANN
// ****************************************************************************************
/*
bool MeanShiftFLANN::init(cv::Mat const &data, cv::Mat const &weights)
{
    if(!MeanShiftStandard::init(data, weights)) return false;

    // flann index
    scaled_data.convertTo(data_float32, CV_32F);
    if(index) delete index;
    index = new cv::flann::Index(data_float32, ip);
    query.resize(D);


    return true;
}


int MeanShiftFLANN::meanshift(cv::Mat const &initial_x, cv::Mat &output_x, Eigen::SparseVector<_Scalar> &output_weights)
{
    _Scalar *ptr_weights = (_Scalar *)weights.data;

    output_x.create(1, D, CV_REAL); // allocate a D-dimensional point

    // variables
	int i, iter, j;
    _Scalar sum_w;
    cv::Mat x = rvect(D);

    // initialisation
	initial_x.copyTo(output_x);
    _Scalar *ptr_output_x = (_Scalar *)output_x.data;

    for(iter = 0; iter < max_iter; ++iter)
    {
        // radius search
        for(i = 0; i < D; ++i) query[i] = ptr_output_x[i]; // creates a query point
        indices.resize(N); dists.resize(N);
        search_params.checks = N; // maximum number of items to search for
        i = index->radiusSearch(query, indices, dists, 3, search_params); // assume radius=3*sigma

        output_weights.resize(i);
        while(--i >= 0)
        {
            output_weights[i].first = j = indices[i];
            output_weights[i].second = exp(-0.5)*dists[i]*dists[i]*ptr_weights[j];
        }

        mean(data_row, output_weights, x, sum_w);

        if(output_weights.empty() || iszero(sum_w))
            break; // can't find any nearby point

        // check if new mean same as old mean
		if(cv::norm(x, output_x, cv::NORM_L1) < 0.000001) break;

        // update new mean
        x.copyTo(output_x);
    }

    return iter;
}
*/


} // namespace MT

#endif
